package net.freebytes.test;

import java.util.*;

class TreeNode {

    int id;
    List<TreeNode> children;
    boolean flag;

    public TreeNode(int id) {
        this.id = id;
        this.children = new ArrayList<>();
    }

    public void addChild(TreeNode child) {
        children.add(child);
    }

    // 深度优先删除未选中的节点
    public boolean dfsRemoveUnselected(Set<Integer> selectedIds) {
        boolean shouldKeep = selectedIds.contains(id); // 当前节点是否被选中
        if (shouldKeep) {
            flag = true;
        }

        // 遍历所有子节点
        Iterator<TreeNode> iterator = children.iterator();
        while (iterator.hasNext()) {
            TreeNode child = iterator.next();
            // 如果子节点及其后代中没有被选中的节点，则删除子节点
            if (!child.dfsRemoveUnselected(selectedIds) ) {
                iterator.remove();
            } else {
                shouldKeep = true; // 如果子节点被选中，当前节点必须保留
            }
        }

        return shouldKeep; // 返回当前节点及其子节点是否需要保留
    }

    // 打印树结构
    public void printTree(String indent) {
        System.out.println(indent + "Node " + id + "                ------------- "+flag);
        for (TreeNode child : children) {
            child.printTree(indent + "  ");
        }
    }
}

public class TreeExample {

    public static void main(String[] args) {
        // 创建树结构
        TreeNode root = new TreeNode(1); // 父节点1
        TreeNode node2 = new TreeNode(2); // 子节点2
        TreeNode node3 = new TreeNode(3); // 父节点3
        TreeNode node4 = new TreeNode(4); // 子节点4
        TreeNode node5 = new TreeNode(5); // 子节点5
        TreeNode node6 = new TreeNode(6); // 子节点6

        // 构建树
        root.addChild(node2);
        root.addChild(node3);
        node3.addChild(node4);
        node3.addChild(node5);
        node5.addChild(node6);
        root.addChild(new TreeNode(7));
        node4.addChild(new TreeNode(41));
        node4.addChild(new TreeNode(42));

        // 模拟选中的节点id集合
        Set<Integer> selectedIds = new HashSet<>(Arrays.asList(41)); // 选择了node4和node6

        System.out.println("Original Tree:");
        root.printTree("");

        // 使用深度优先的方式删除未选中的节点
        root.dfsRemoveUnselected(selectedIds);

        System.out.println("\nTree After Removing Unselected Nodes:");
        root.printTree("");
    }
}
